{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88dc24ff",
   "metadata": {},
   "source": [
    "## 用卷积神经网络训练Cifar100\n",
    "For this tutorial, we will use the CIFAR100 dataset. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c2411fc",
   "metadata": {},
   "source": [
    "## 1. Load and normalize CIFAR100\n",
    "Using torchvision, it's extremely easy to load CIFAR100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f7963711",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.optim.lr_scheduler as lr_scheduler\n",
    "# import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ed493286",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using PyTorch version: 1.9.1+cu111  Device: cuda\n"
     ]
    }
   ],
   "source": [
    "# Set the device\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "    torch.cuda.set_device(0)\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "print('Using PyTorch version:', torch.__version__, ' Device:', device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12bc2962",
   "metadata": {},
   "source": [
    "The output of torchvision datasets are PILImage images of range [0, 1]. We transform them to Tensors of normalized range [-1, 1]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2dcf4ab3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "# Define the batch size and number of epochs\n",
    "batch_size = 128\n",
    "num_epochs = 10\n",
    "\n",
    "# Define the transformation for data preprocessing\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalize the input\n",
    "])\n",
    "\n",
    "# Load the CIFAR-100 dataset\n",
    "train_dataset = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=transform)\n",
    "test_dataset = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform)\n",
    "\n",
    "# Create data loaders\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07f934de",
   "metadata": {},
   "source": [
    "## 2. Define SEK\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "701a41c7-3b41-429b-99d7-c56335400ca6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "import torch.nn.functional as F\n",
    "\n",
    "__all__ = ['mbv2_ca']\n",
    "\n",
    "class ConvBNReLU(nn.Sequential):\n",
    "    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):\n",
    "        padding = (kernel_size - 1) // 2\n",
    "        super(ConvBNReLU, self).__init__(\n",
    "            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),\n",
    "            nn.BatchNorm2d(out_planes),\n",
    "            nn.ReLU6(inplace=True)\n",
    "        )\n",
    "\n",
    "class h_sigmoid(nn.Module):\n",
    "    def __init__(self, inplace=True):\n",
    "        super(h_sigmoid, self).__init__()\n",
    "        self.relu = nn.ReLU6(inplace=inplace)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.relu(x + 3) / 6\n",
    "\n",
    "class h_swish(nn.Module):\n",
    "    def __init__(self, inplace=True):\n",
    "        super(h_swish, self).__init__()\n",
    "        self.sigmoid = h_sigmoid(inplace=inplace)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x * self.sigmoid(x)\n",
    "\n",
    "class swish(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return x * torch.sigmoid(x)\n",
    "    \n",
    "class CoordAtt(nn.Module):\n",
    "    def __init__(self, inp, oup, groups=32):\n",
    "        super(CoordAtt, self).__init__()\n",
    "        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))\n",
    "        self.pool_w = nn.AdaptiveAvgPool2d((1, None))\n",
    "\n",
    "        mip = max(8, inp // groups)\n",
    "\n",
    "        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)\n",
    "        self.bn1 = nn.BatchNorm2d(mip)\n",
    "        self.conv2 = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)\n",
    "        self.conv3 = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)\n",
    "        self.relu = h_swish()\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "        n,c,h,w = x.size()\n",
    "        x_h = self.pool_h(x)\n",
    "        x_w = self.pool_w(x).permute(0, 1, 3, 2)\n",
    "\n",
    "        y = torch.cat([x_h, x_w], dim=2)\n",
    "        y = self.conv1(y)\n",
    "        y = self.bn1(y)\n",
    "        y = self.relu(y) \n",
    "        x_h, x_w = torch.split(y, [h, w], dim=2)\n",
    "        x_w = x_w.permute(0, 1, 3, 2)\n",
    "\n",
    "        x_h = self.conv2(x_h).sigmoid()\n",
    "        x_w = self.conv3(x_w).sigmoid()\n",
    "        x_h = x_h.expand(-1, -1, h, w)\n",
    "        x_w = x_w.expand(-1, -1, h, w)\n",
    "\n",
    "        y = identity * x_w * x_h\n",
    "\n",
    "        return y\n",
    "\n",
    "def _make_divisible(v, divisor, min_value=None):\n",
    "    \"\"\"\n",
    "    This function is taken from the original tf repo.\n",
    "    It ensures that all layers have a channel number that is divisible by 8\n",
    "    It can be seen here:\n",
    "    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py\n",
    "    :param v:\n",
    "    :param divisor:\n",
    "    :param min_value:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    if min_value is None:\n",
    "        min_value = divisor\n",
    "    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)\n",
    "    # Make sure that round down does not go down by more than 10%.\n",
    "    if new_v < 0.9 * v:\n",
    "        new_v += divisor\n",
    "    return new_v\n",
    "\n",
    "\n",
    "def conv_3x3_bn(inp, oup, stride):\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),\n",
    "        nn.BatchNorm2d(oup),\n",
    "        nn.ReLU6(inplace=True)\n",
    "    )\n",
    "\n",
    "\n",
    "def conv_1x1_bn(inp, oup):\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),\n",
    "        nn.BatchNorm2d(oup),\n",
    "        nn.ReLU6(inplace=True)\n",
    "    )\n",
    "\n",
    "class InvertedResidual(nn.Module):\n",
    "    def __init__(self, inp, oup, stride, expand_ratio):\n",
    "        super(InvertedResidual, self).__init__()\n",
    "        assert stride in [1, 2]\n",
    "\n",
    "        hidden_dim = round(inp * expand_ratio)\n",
    "        self.identity = stride == 1 and inp == oup\n",
    "\n",
    "        if expand_ratio == 1:\n",
    "            self.conv = nn.Sequential(\n",
    "                # dw\n",
    "                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),\n",
    "                nn.BatchNorm2d(hidden_dim),\n",
    "                nn.ReLU6(inplace=True),\n",
    "                # pw-linear\n",
    "                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),\n",
    "                nn.BatchNorm2d(oup),\n",
    "            )\n",
    "        else:\n",
    "            self.conv = nn.Sequential(\n",
    "                # pw\n",
    "                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),\n",
    "                nn.BatchNorm2d(hidden_dim),\n",
    "                nn.ReLU6(inplace=True),\n",
    "                # dw\n",
    "                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),\n",
    "                nn.BatchNorm2d(hidden_dim),\n",
    "                nn.ReLU6(inplace=True),\n",
    "                # coordinate attention\n",
    "                CoordAtt(hidden_dim, hidden_dim),\n",
    "                # pw-linear\n",
    "                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),\n",
    "                nn.BatchNorm2d(oup),\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = self.conv(x)\n",
    "        if self.identity:\n",
    "            return x + y\n",
    "        else:\n",
    "            return y\n",
    "\n",
    "class MBV2_CA(nn.Module):\n",
    "    def __init__(self, num_classes=100, width_mult=1.):\n",
    "        super(MBV2_CA, self).__init__()\n",
    "        # setting of inverted residual blocks\n",
    "        self.cfgs = [\n",
    "            # t, c, n, s\n",
    "            [1,  16, 1, 1],\n",
    "            [6,  24, 2, 2],\n",
    "            [6,  32, 3, 2],\n",
    "            [6,  64, 4, 2],\n",
    "            [6,  96, 3, 1],\n",
    "            [6, 160, 3, 2],\n",
    "            [6, 320, 1, 1],\n",
    "        ]\n",
    "\n",
    "        # building first layer\n",
    "        input_channel = _make_divisible(32 * width_mult, 4 if width_mult == 0.1 else 8)\n",
    "        layers = [conv_3x3_bn(3, input_channel, 2)]\n",
    "        # building inverted residual blocks\n",
    "        block = InvertedResidual\n",
    "        for t, c, n, s in self.cfgs:\n",
    "            output_channel = _make_divisible(c * width_mult, 4 if width_mult == 0.1 else 8)\n",
    "            for i in range(n):\n",
    "                layers.append(block(input_channel, output_channel, s if i == 0 else 1, t))\n",
    "                input_channel = output_channel\n",
    "        self.features = nn.Sequential(*layers)\n",
    "        # building last several layers\n",
    "        output_channel = _make_divisible(1280 * width_mult, 4 if width_mult == 0.1 else 8) if width_mult > 1.0 else 1280\n",
    "        self.conv = conv_1x1_bn(input_channel, output_channel)\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.classifier = nn.Sequential(\n",
    "                nn.Dropout(0.1),\n",
    "                nn.Linear(output_channel, num_classes)\n",
    "                )\n",
    "\n",
    "        self._initialize_weights()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = self.conv(x)\n",
    "        x = self.avgpool(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.classifier(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def _initialize_weights(self):\n",
    "        for m in self.modules():\n",
    "            #print(m)\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                #print(m.weight.size())\n",
    "                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
    "                m.weight.data.normal_(0, math.sqrt(2. / n))\n",
    "                if m.bias is not None:\n",
    "                    m.bias.data.zero_()\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                m.weight.data.fill_(1)\n",
    "                m.bias.data.zero_()\n",
    "            elif isinstance(m, nn.Linear):\n",
    "                m.weight.data.normal_(0, 0.01)\n",
    "                m.bias.data.zero_()\n",
    "\n",
    "\n",
    "def mbv2_ca(**kwargs):\n",
    "    return  MBV2_CA(**kwargs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c2dbf7f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the model\n",
    "model = mbv2_ca().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a34acf3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model\n",
    "# PATH = './cifar_skn3.pth'\n",
    "# model.load_state_dict(torch.load(PATH))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12312f6f",
   "metadata": {},
   "source": [
    "## 3. Define a Loss function and optimizer\n",
    "Let’s use a Classification Cross-Entropy loss and Adam."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "40fbef67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.05, momentum=0.9, weight_decay=5e-4)\n",
    "\n",
    "# Define the learning rate schedule\n",
    "lr_schedule_milestones = [int(num_epochs * 0.5), int(num_epochs * 0.75)]  # 50% and 75% of total training epochs\n",
    "lr_schedule_gamma = 0.1  # Learning rate reduction factor\n",
    "\n",
    "# Define the learning rate scheduler\n",
    "scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=lr_schedule_milestones, gamma=lr_schedule_gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ceb4d7b",
   "metadata": {},
   "source": [
    "## 4. Train the network\n",
    "This is when things start to get interesting. We simply have to loop over our data iterator, and feed the inputs to the network and optimize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "654c8151",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lists to store training and testing losses, and accuracies\n",
    "train_losses = []\n",
    "test_losses = []\n",
    "train_accs = []\n",
    "test_accs = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3f8b32e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/10], Step [10/391], Loss: 4.1110\n",
      "Epoch [1/10], Step [20/391], Loss: 4.0972\n",
      "Epoch [1/10], Step [30/391], Loss: 4.1053\n",
      "Epoch [1/10], Step [40/391], Loss: 4.0860\n",
      "Epoch [1/10], Step [50/391], Loss: 4.0708\n",
      "Epoch [1/10], Step [60/391], Loss: 4.0717\n",
      "Epoch [1/10], Step [70/391], Loss: 4.0676\n",
      "Epoch [1/10], Step [80/391], Loss: 4.0635\n",
      "Epoch [1/10], Step [90/391], Loss: 4.0529\n",
      "Epoch [1/10], Step [100/391], Loss: 4.0483\n",
      "Epoch [1/10], Step [110/391], Loss: 4.0445\n",
      "Epoch [1/10], Step [120/391], Loss: 4.0446\n",
      "Epoch [1/10], Step [130/391], Loss: 4.0430\n",
      "Epoch [1/10], Step [140/391], Loss: 4.0549\n",
      "Epoch [1/10], Step [150/391], Loss: 4.0616\n",
      "Epoch [1/10], Step [160/391], Loss: 4.0661\n",
      "Epoch [1/10], Step [170/391], Loss: 4.0610\n",
      "Epoch [1/10], Step [180/391], Loss: 4.0613\n",
      "Epoch [1/10], Step [190/391], Loss: 4.0537\n",
      "Epoch [1/10], Step [200/391], Loss: 4.0494\n",
      "Epoch [1/10], Step [210/391], Loss: 4.0451\n",
      "Epoch [1/10], Step [220/391], Loss: 4.0425\n",
      "Epoch [1/10], Step [230/391], Loss: 4.0386\n",
      "Epoch [1/10], Step [240/391], Loss: 4.0322\n",
      "Epoch [1/10], Step [250/391], Loss: 4.0279\n",
      "Epoch [1/10], Step [260/391], Loss: 4.0230\n",
      "Epoch [1/10], Step [270/391], Loss: 4.0175\n",
      "Epoch [1/10], Step [280/391], Loss: 4.0107\n",
      "Epoch [1/10], Step [290/391], Loss: 4.0101\n",
      "Epoch [1/10], Step [300/391], Loss: 4.0146\n",
      "Epoch [1/10], Step [310/391], Loss: 4.0170\n",
      "Epoch [1/10], Step [320/391], Loss: 4.0176\n",
      "Epoch [1/10], Step [330/391], Loss: 4.0156\n",
      "Epoch [1/10], Step [340/391], Loss: 4.0127\n",
      "Epoch [1/10], Step [350/391], Loss: 4.0093\n",
      "Epoch [1/10], Step [360/391], Loss: 4.0066\n",
      "Epoch [1/10], Step [370/391], Loss: 4.0041\n",
      "Epoch [1/10], Step [380/391], Loss: 4.0001\n",
      "Epoch [1/10], Step [390/391], Loss: 3.9972\n",
      "Epoch [1/10], Test Accuracy: 9.08%\n",
      "Epoch [2/10], Step [10/391], Loss: 3.8609\n",
      "Epoch [2/10], Step [20/391], Loss: 3.8654\n",
      "Epoch [2/10], Step [30/391], Loss: 3.8427\n",
      "Epoch [2/10], Step [40/391], Loss: 3.8284\n",
      "Epoch [2/10], Step [50/391], Loss: 3.8175\n",
      "Epoch [2/10], Step [60/391], Loss: 3.8162\n",
      "Epoch [2/10], Step [70/391], Loss: 3.8101\n",
      "Epoch [2/10], Step [80/391], Loss: 3.8013\n",
      "Epoch [2/10], Step [90/391], Loss: 3.8011\n",
      "Epoch [2/10], Step [100/391], Loss: 3.7941\n",
      "Epoch [2/10], Step [110/391], Loss: 3.7935\n",
      "Epoch [2/10], Step [120/391], Loss: 3.7877\n",
      "Epoch [2/10], Step [130/391], Loss: 3.7875\n",
      "Epoch [2/10], Step [140/391], Loss: 3.7845\n",
      "Epoch [2/10], Step [150/391], Loss: 3.7813\n",
      "Epoch [2/10], Step [160/391], Loss: 3.7754\n",
      "Epoch [2/10], Step [170/391], Loss: 3.7718\n",
      "Epoch [2/10], Step [180/391], Loss: 3.7663\n",
      "Epoch [2/10], Step [190/391], Loss: 3.7601\n",
      "Epoch [2/10], Step [200/391], Loss: 3.7575\n",
      "Epoch [2/10], Step [210/391], Loss: 3.7521\n",
      "Epoch [2/10], Step [220/391], Loss: 3.7464\n",
      "Epoch [2/10], Step [230/391], Loss: 3.7413\n",
      "Epoch [2/10], Step [240/391], Loss: 3.7354\n",
      "Epoch [2/10], Step [250/391], Loss: 3.7319\n",
      "Epoch [2/10], Step [260/391], Loss: 3.7288\n",
      "Epoch [2/10], Step [270/391], Loss: 3.7270\n",
      "Epoch [2/10], Step [280/391], Loss: 3.7233\n",
      "Epoch [2/10], Step [290/391], Loss: 3.7195\n",
      "Epoch [2/10], Step [300/391], Loss: 3.7171\n",
      "Epoch [2/10], Step [310/391], Loss: 3.7130\n",
      "Epoch [2/10], Step [320/391], Loss: 3.7092\n",
      "Epoch [2/10], Step [330/391], Loss: 3.7064\n",
      "Epoch [2/10], Step [340/391], Loss: 3.7043\n",
      "Epoch [2/10], Step [350/391], Loss: 3.6984\n",
      "Epoch [2/10], Step [360/391], Loss: 3.6947\n",
      "Epoch [2/10], Step [370/391], Loss: 3.6909\n",
      "Epoch [2/10], Step [380/391], Loss: 3.6880\n",
      "Epoch [2/10], Step [390/391], Loss: 3.6855\n",
      "Epoch [2/10], Test Accuracy: 14.83%\n",
      "Epoch [3/10], Step [10/391], Loss: 3.4406\n",
      "Epoch [3/10], Step [20/391], Loss: 3.4696\n",
      "Epoch [3/10], Step [30/391], Loss: 3.4782\n",
      "Epoch [3/10], Step [40/391], Loss: 3.4807\n",
      "Epoch [3/10], Step [50/391], Loss: 3.4784\n",
      "Epoch [3/10], Step [60/391], Loss: 3.4809\n",
      "Epoch [3/10], Step [70/391], Loss: 3.4763\n",
      "Epoch [3/10], Step [80/391], Loss: 3.4896\n",
      "Epoch [3/10], Step [90/391], Loss: 3.4991\n",
      "Epoch [3/10], Step [100/391], Loss: 3.5000\n",
      "Epoch [3/10], Step [110/391], Loss: 3.4959\n",
      "Epoch [3/10], Step [120/391], Loss: 3.4953\n",
      "Epoch [3/10], Step [130/391], Loss: 3.4929\n",
      "Epoch [3/10], Step [140/391], Loss: 3.4923\n",
      "Epoch [3/10], Step [150/391], Loss: 3.4902\n",
      "Epoch [3/10], Step [160/391], Loss: 3.4908\n",
      "Epoch [3/10], Step [170/391], Loss: 3.4861\n",
      "Epoch [3/10], Step [180/391], Loss: 3.4843\n",
      "Epoch [3/10], Step [190/391], Loss: 3.4798\n",
      "Epoch [3/10], Step [200/391], Loss: 3.4794\n",
      "Epoch [3/10], Step [210/391], Loss: 3.4760\n",
      "Epoch [3/10], Step [220/391], Loss: 3.4732\n",
      "Epoch [3/10], Step [230/391], Loss: 3.4704\n",
      "Epoch [3/10], Step [240/391], Loss: 3.4701\n",
      "Epoch [3/10], Step [250/391], Loss: 3.4682\n",
      "Epoch [3/10], Step [260/391], Loss: 3.4655\n",
      "Epoch [3/10], Step [270/391], Loss: 3.4628\n",
      "Epoch [3/10], Step [280/391], Loss: 3.4593\n",
      "Epoch [3/10], Step [290/391], Loss: 3.4561\n",
      "Epoch [3/10], Step [300/391], Loss: 3.4539\n",
      "Epoch [3/10], Step [310/391], Loss: 3.4517\n",
      "Epoch [3/10], Step [320/391], Loss: 3.4492\n",
      "Epoch [3/10], Step [330/391], Loss: 3.4489\n",
      "Epoch [3/10], Step [340/391], Loss: 3.4478\n",
      "Epoch [3/10], Step [350/391], Loss: 3.4450\n",
      "Epoch [3/10], Step [360/391], Loss: 3.4431\n",
      "Epoch [3/10], Step [370/391], Loss: 3.4392\n",
      "Epoch [3/10], Step [380/391], Loss: 3.4364\n",
      "Epoch [3/10], Step [390/391], Loss: 3.4339\n",
      "Epoch [3/10], Test Accuracy: 17.78%\n",
      "Epoch [4/10], Step [10/391], Loss: 3.2977\n",
      "Epoch [4/10], Step [20/391], Loss: 3.2828\n",
      "Epoch [4/10], Step [30/391], Loss: 3.2466\n",
      "Epoch [4/10], Step [40/391], Loss: 3.2661\n",
      "Epoch [4/10], Step [50/391], Loss: 3.2719\n",
      "Epoch [4/10], Step [60/391], Loss: 3.2777\n",
      "Epoch [4/10], Step [70/391], Loss: 3.2780\n",
      "Epoch [4/10], Step [80/391], Loss: 3.2794\n",
      "Epoch [4/10], Step [90/391], Loss: 3.2803\n",
      "Epoch [4/10], Step [100/391], Loss: 3.2789\n",
      "Epoch [4/10], Step [110/391], Loss: 3.2836\n",
      "Epoch [4/10], Step [120/391], Loss: 3.2817\n",
      "Epoch [4/10], Step [130/391], Loss: 3.2812\n",
      "Epoch [4/10], Step [140/391], Loss: 3.2765\n",
      "Epoch [4/10], Step [150/391], Loss: 3.2793\n",
      "Epoch [4/10], Step [160/391], Loss: 3.2784\n",
      "Epoch [4/10], Step [170/391], Loss: 3.2837\n",
      "Epoch [4/10], Step [180/391], Loss: 3.2889\n",
      "Epoch [4/10], Step [190/391], Loss: 3.2896\n",
      "Epoch [4/10], Step [200/391], Loss: 3.2880\n",
      "Epoch [4/10], Step [210/391], Loss: 3.2860\n",
      "Epoch [4/10], Step [220/391], Loss: 3.2853\n",
      "Epoch [4/10], Step [230/391], Loss: 3.2869\n",
      "Epoch [4/10], Step [240/391], Loss: 3.2877\n",
      "Epoch [4/10], Step [250/391], Loss: 3.2909\n",
      "Epoch [4/10], Step [260/391], Loss: 3.2942\n",
      "Epoch [4/10], Step [270/391], Loss: 3.2991\n",
      "Epoch [4/10], Step [280/391], Loss: 3.2985\n",
      "Epoch [4/10], Step [290/391], Loss: 3.2977\n",
      "Epoch [4/10], Step [300/391], Loss: 3.2975\n",
      "Epoch [4/10], Step [310/391], Loss: 3.2963\n",
      "Epoch [4/10], Step [320/391], Loss: 3.2957\n",
      "Epoch [4/10], Step [330/391], Loss: 3.2929\n",
      "Epoch [4/10], Step [340/391], Loss: 3.2936\n",
      "Epoch [4/10], Step [350/391], Loss: 3.2925\n",
      "Epoch [4/10], Step [360/391], Loss: 3.2948\n",
      "Epoch [4/10], Step [370/391], Loss: 3.2956\n",
      "Epoch [4/10], Step [380/391], Loss: 3.2949\n",
      "Epoch [4/10], Step [390/391], Loss: 3.2935\n",
      "Epoch [4/10], Test Accuracy: 20.37%\n",
      "Epoch [5/10], Step [10/391], Loss: 3.0935\n",
      "Epoch [5/10], Step [20/391], Loss: 3.1420\n",
      "Epoch [5/10], Step [30/391], Loss: 3.1465\n",
      "Epoch [5/10], Step [40/391], Loss: 3.1612\n",
      "Epoch [5/10], Step [50/391], Loss: 3.1553\n",
      "Epoch [5/10], Step [60/391], Loss: 3.1536\n",
      "Epoch [5/10], Step [70/391], Loss: 3.1629\n",
      "Epoch [5/10], Step [80/391], Loss: 3.1560\n",
      "Epoch [5/10], Step [90/391], Loss: 3.1495\n",
      "Epoch [5/10], Step [100/391], Loss: 3.1597\n",
      "Epoch [5/10], Step [110/391], Loss: 3.1634\n",
      "Epoch [5/10], Step [120/391], Loss: 3.1644\n",
      "Epoch [5/10], Step [130/391], Loss: 3.1632\n",
      "Epoch [5/10], Step [140/391], Loss: 3.1685\n",
      "Epoch [5/10], Step [150/391], Loss: 3.1695\n",
      "Epoch [5/10], Step [160/391], Loss: 3.1622\n",
      "Epoch [5/10], Step [170/391], Loss: 3.1653\n",
      "Epoch [5/10], Step [180/391], Loss: 3.1680\n",
      "Epoch [5/10], Step [190/391], Loss: 3.1682\n",
      "Epoch [5/10], Step [200/391], Loss: 3.1743\n",
      "Epoch [5/10], Step [210/391], Loss: 3.1770\n",
      "Epoch [5/10], Step [220/391], Loss: 3.1823\n",
      "Epoch [5/10], Step [230/391], Loss: 3.1839\n",
      "Epoch [5/10], Step [240/391], Loss: 3.1819\n",
      "Epoch [5/10], Step [250/391], Loss: 3.1789\n",
      "Epoch [5/10], Step [260/391], Loss: 3.1777\n",
      "Epoch [5/10], Step [270/391], Loss: 3.1770\n",
      "Epoch [5/10], Step [280/391], Loss: 3.1762\n",
      "Epoch [5/10], Step [290/391], Loss: 3.1734\n",
      "Epoch [5/10], Step [300/391], Loss: 3.1693\n",
      "Epoch [5/10], Step [310/391], Loss: 3.1676\n",
      "Epoch [5/10], Step [320/391], Loss: 3.1658\n",
      "Epoch [5/10], Step [330/391], Loss: 3.1649\n",
      "Epoch [5/10], Step [340/391], Loss: 3.1649\n",
      "Epoch [5/10], Step [350/391], Loss: 3.1656\n",
      "Epoch [5/10], Step [360/391], Loss: 3.1660\n",
      "Epoch [5/10], Step [370/391], Loss: 3.1661\n",
      "Epoch [5/10], Step [380/391], Loss: 3.1656\n",
      "Epoch [5/10], Step [390/391], Loss: 3.1654\n",
      "Epoch [5/10], Test Accuracy: 21.20%\n",
      "Epoch [6/10], Step [10/391], Loss: 3.0431\n",
      "Epoch [6/10], Step [20/391], Loss: 3.0636\n",
      "Epoch [6/10], Step [30/391], Loss: 3.0537\n",
      "Epoch [6/10], Step [40/391], Loss: 3.0470\n",
      "Epoch [6/10], Step [50/391], Loss: 3.0447\n",
      "Epoch [6/10], Step [60/391], Loss: 3.0744\n",
      "Epoch [6/10], Step [70/391], Loss: 3.0900\n",
      "Epoch [6/10], Step [80/391], Loss: 3.0866\n",
      "Epoch [6/10], Step [90/391], Loss: 3.0814\n",
      "Epoch [6/10], Step [100/391], Loss: 3.0749\n",
      "Epoch [6/10], Step [110/391], Loss: 3.0755\n",
      "Epoch [6/10], Step [120/391], Loss: 3.0747\n",
      "Epoch [6/10], Step [130/391], Loss: 3.0771\n",
      "Epoch [6/10], Step [140/391], Loss: 3.0700\n",
      "Epoch [6/10], Step [150/391], Loss: 3.0796\n",
      "Epoch [6/10], Step [160/391], Loss: 3.0795\n",
      "Epoch [6/10], Step [170/391], Loss: 3.0815\n",
      "Epoch [6/10], Step [180/391], Loss: 3.0808\n",
      "Epoch [6/10], Step [190/391], Loss: 3.0862\n",
      "Epoch [6/10], Step [200/391], Loss: 3.0950\n",
      "Epoch [6/10], Step [210/391], Loss: 3.0994\n",
      "Epoch [6/10], Step [220/391], Loss: 3.1029\n",
      "Epoch [6/10], Step [230/391], Loss: 3.1086\n",
      "Epoch [6/10], Step [240/391], Loss: 3.1108\n",
      "Epoch [6/10], Step [250/391], Loss: 3.1070\n",
      "Epoch [6/10], Step [260/391], Loss: 3.1047\n",
      "Epoch [6/10], Step [270/391], Loss: 3.1021\n",
      "Epoch [6/10], Step [280/391], Loss: 3.1002\n",
      "Epoch [6/10], Step [290/391], Loss: 3.1025\n",
      "Epoch [6/10], Step [300/391], Loss: 3.0991\n",
      "Epoch [6/10], Step [310/391], Loss: 3.0972\n",
      "Epoch [6/10], Step [320/391], Loss: 3.0941\n",
      "Epoch [6/10], Step [330/391], Loss: 3.0901\n",
      "Epoch [6/10], Step [340/391], Loss: 3.0883\n",
      "Epoch [6/10], Step [350/391], Loss: 3.0856\n",
      "Epoch [6/10], Step [360/391], Loss: 3.0832\n",
      "Epoch [6/10], Step [370/391], Loss: 3.0807\n",
      "Epoch [6/10], Step [380/391], Loss: 3.0797\n",
      "Epoch [6/10], Step [390/391], Loss: 3.0825\n",
      "Epoch [6/10], Test Accuracy: 14.32%\n",
      "Epoch [7/10], Step [10/391], Loss: 3.1693\n",
      "Epoch [7/10], Step [20/391], Loss: 3.1131\n",
      "Epoch [7/10], Step [30/391], Loss: 3.0934\n",
      "Epoch [7/10], Step [40/391], Loss: 3.0781\n",
      "Epoch [7/10], Step [50/391], Loss: 3.0862\n",
      "Epoch [7/10], Step [60/391], Loss: 3.0720\n",
      "Epoch [7/10], Step [70/391], Loss: 3.0699\n",
      "Epoch [7/10], Step [80/391], Loss: 3.0640\n",
      "Epoch [7/10], Step [90/391], Loss: 3.0549\n",
      "Epoch [7/10], Step [100/391], Loss: 3.0479\n",
      "Epoch [7/10], Step [110/391], Loss: 3.0395\n",
      "Epoch [7/10], Step [120/391], Loss: 3.0361\n",
      "Epoch [7/10], Step [130/391], Loss: 3.0344\n",
      "Epoch [7/10], Step [140/391], Loss: 3.0319\n",
      "Epoch [7/10], Step [150/391], Loss: 3.0229\n",
      "Epoch [7/10], Step [160/391], Loss: 3.0174\n",
      "Epoch [7/10], Step [170/391], Loss: 3.0095\n",
      "Epoch [7/10], Step [180/391], Loss: 3.0064\n",
      "Epoch [7/10], Step [190/391], Loss: 3.0059\n",
      "Epoch [7/10], Step [200/391], Loss: 3.0067\n",
      "Epoch [7/10], Step [210/391], Loss: 3.0034\n",
      "Epoch [7/10], Step [220/391], Loss: 3.0037\n",
      "Epoch [7/10], Step [230/391], Loss: 3.0006\n",
      "Epoch [7/10], Step [240/391], Loss: 3.0008\n",
      "Epoch [7/10], Step [250/391], Loss: 2.9986\n",
      "Epoch [7/10], Step [260/391], Loss: 2.9955\n",
      "Epoch [7/10], Step [270/391], Loss: 2.9951\n",
      "Epoch [7/10], Step [280/391], Loss: 2.9912\n",
      "Epoch [7/10], Step [290/391], Loss: 2.9881\n",
      "Epoch [7/10], Step [300/391], Loss: 2.9861\n",
      "Epoch [7/10], Step [310/391], Loss: 2.9851\n",
      "Epoch [7/10], Step [320/391], Loss: 2.9848\n",
      "Epoch [7/10], Step [330/391], Loss: 2.9862\n",
      "Epoch [7/10], Step [340/391], Loss: 2.9848\n",
      "Epoch [7/10], Step [350/391], Loss: 2.9860\n",
      "Epoch [7/10], Step [360/391], Loss: 2.9841\n",
      "Epoch [7/10], Step [370/391], Loss: 2.9843\n",
      "Epoch [7/10], Step [380/391], Loss: 2.9842\n",
      "Epoch [7/10], Step [390/391], Loss: 2.9812\n",
      "Epoch [7/10], Test Accuracy: 26.12%\n",
      "Epoch [8/10], Step [10/391], Loss: 2.7123\n",
      "Epoch [8/10], Step [20/391], Loss: 2.7490\n",
      "Epoch [8/10], Step [30/391], Loss: 2.7651\n",
      "Epoch [8/10], Step [40/391], Loss: 2.7806\n",
      "Epoch [8/10], Step [50/391], Loss: 2.7966\n",
      "Epoch [8/10], Step [60/391], Loss: 2.8103\n",
      "Epoch [8/10], Step [70/391], Loss: 2.8188\n",
      "Epoch [8/10], Step [80/391], Loss: 2.8061\n",
      "Epoch [8/10], Step [90/391], Loss: 2.8111\n",
      "Epoch [8/10], Step [100/391], Loss: 2.8066\n",
      "Epoch [8/10], Step [110/391], Loss: 2.8106\n",
      "Epoch [8/10], Step [120/391], Loss: 2.8145\n",
      "Epoch [8/10], Step [130/391], Loss: 2.8176\n",
      "Epoch [8/10], Step [140/391], Loss: 2.8187\n",
      "Epoch [8/10], Step [150/391], Loss: 2.8170\n",
      "Epoch [8/10], Step [160/391], Loss: 2.8139\n",
      "Epoch [8/10], Step [170/391], Loss: 2.8138\n",
      "Epoch [8/10], Step [180/391], Loss: 2.8183\n",
      "Epoch [8/10], Step [190/391], Loss: 2.8167\n",
      "Epoch [8/10], Step [200/391], Loss: 2.8197\n",
      "Epoch [8/10], Step [210/391], Loss: 2.8198\n",
      "Epoch [8/10], Step [220/391], Loss: 2.8221\n",
      "Epoch [8/10], Step [230/391], Loss: 2.8238\n",
      "Epoch [8/10], Step [240/391], Loss: 2.8269\n",
      "Epoch [8/10], Step [250/391], Loss: 2.8262\n",
      "Epoch [8/10], Step [260/391], Loss: 2.8295\n",
      "Epoch [8/10], Step [270/391], Loss: 2.8306\n",
      "Epoch [8/10], Step [280/391], Loss: 2.8324\n",
      "Epoch [8/10], Step [290/391], Loss: 2.8359\n",
      "Epoch [8/10], Step [300/391], Loss: 2.8379\n",
      "Epoch [8/10], Step [310/391], Loss: 2.8384\n",
      "Epoch [8/10], Step [320/391], Loss: 2.8377\n",
      "Epoch [8/10], Step [330/391], Loss: 2.8398\n",
      "Epoch [8/10], Step [340/391], Loss: 2.8381\n",
      "Epoch [8/10], Step [350/391], Loss: 2.8370\n",
      "Epoch [8/10], Step [360/391], Loss: 2.8373\n",
      "Epoch [8/10], Step [370/391], Loss: 2.8345\n",
      "Epoch [8/10], Step [380/391], Loss: 2.8333\n",
      "Epoch [8/10], Step [390/391], Loss: 2.8324\n",
      "Epoch [8/10], Test Accuracy: 27.50%\n",
      "Epoch [9/10], Step [10/391], Loss: 2.6018\n",
      "Epoch [9/10], Step [20/391], Loss: 2.6597\n",
      "Epoch [9/10], Step [30/391], Loss: 2.6703\n",
      "Epoch [9/10], Step [40/391], Loss: 2.6655\n",
      "Epoch [9/10], Step [50/391], Loss: 2.6780\n",
      "Epoch [9/10], Step [60/391], Loss: 2.6761\n",
      "Epoch [9/10], Step [70/391], Loss: 2.6825\n",
      "Epoch [9/10], Step [80/391], Loss: 2.6821\n",
      "Epoch [9/10], Step [90/391], Loss: 2.6823\n",
      "Epoch [9/10], Step [100/391], Loss: 2.6895\n",
      "Epoch [9/10], Step [110/391], Loss: 2.7075\n",
      "Epoch [9/10], Step [120/391], Loss: 2.7362\n",
      "Epoch [9/10], Step [130/391], Loss: 2.7585\n",
      "Epoch [9/10], Step [140/391], Loss: 2.7640\n",
      "Epoch [9/10], Step [150/391], Loss: 2.7651\n",
      "Epoch [9/10], Step [160/391], Loss: 2.7779\n",
      "Epoch [9/10], Step [170/391], Loss: 2.7784\n",
      "Epoch [9/10], Step [180/391], Loss: 2.7814\n",
      "Epoch [9/10], Step [190/391], Loss: 2.7784\n",
      "Epoch [9/10], Step [200/391], Loss: 2.7772\n",
      "Epoch [9/10], Step [210/391], Loss: 2.7775\n",
      "Epoch [9/10], Step [220/391], Loss: 2.7741\n",
      "Epoch [9/10], Step [230/391], Loss: 2.7752\n",
      "Epoch [9/10], Step [240/391], Loss: 2.7775\n",
      "Epoch [9/10], Step [250/391], Loss: 2.7779\n",
      "Epoch [9/10], Step [260/391], Loss: 2.7788\n",
      "Epoch [9/10], Step [270/391], Loss: 2.7786\n",
      "Epoch [9/10], Step [280/391], Loss: 2.7828\n",
      "Epoch [9/10], Step [290/391], Loss: 2.7889\n",
      "Epoch [9/10], Step [300/391], Loss: 2.7899\n",
      "Epoch [9/10], Step [310/391], Loss: 2.7931\n",
      "Epoch [9/10], Step [320/391], Loss: 2.7940\n",
      "Epoch [9/10], Step [330/391], Loss: 2.7965\n",
      "Epoch [9/10], Step [340/391], Loss: 2.7971\n",
      "Epoch [9/10], Step [350/391], Loss: 2.7986\n",
      "Epoch [9/10], Step [360/391], Loss: 2.7978\n",
      "Epoch [9/10], Step [370/391], Loss: 2.7948\n",
      "Epoch [9/10], Step [380/391], Loss: 2.7944\n",
      "Epoch [9/10], Step [390/391], Loss: 2.7927\n",
      "Epoch [9/10], Test Accuracy: 28.57%\n",
      "Epoch [10/10], Step [10/391], Loss: 2.6031\n",
      "Epoch [10/10], Step [20/391], Loss: 2.6099\n",
      "Epoch [10/10], Step [30/391], Loss: 2.6273\n",
      "Epoch [10/10], Step [40/391], Loss: 2.6339\n",
      "Epoch [10/10], Step [50/391], Loss: 2.6245\n",
      "Epoch [10/10], Step [60/391], Loss: 2.6412\n",
      "Epoch [10/10], Step [70/391], Loss: 2.6260\n",
      "Epoch [10/10], Step [80/391], Loss: 2.6335\n",
      "Epoch [10/10], Step [90/391], Loss: 2.6239\n",
      "Epoch [10/10], Step [100/391], Loss: 2.6276\n",
      "Epoch [10/10], Step [110/391], Loss: 2.6323\n",
      "Epoch [10/10], Step [120/391], Loss: 2.6393\n",
      "Epoch [10/10], Step [130/391], Loss: 2.6399\n",
      "Epoch [10/10], Step [140/391], Loss: 2.6437\n",
      "Epoch [10/10], Step [150/391], Loss: 2.6516\n",
      "Epoch [10/10], Step [160/391], Loss: 2.6497\n",
      "Epoch [10/10], Step [170/391], Loss: 2.6472\n",
      "Epoch [10/10], Step [180/391], Loss: 2.6474\n",
      "Epoch [10/10], Step [190/391], Loss: 2.6484\n",
      "Epoch [10/10], Step [200/391], Loss: 2.6536\n",
      "Epoch [10/10], Step [210/391], Loss: 2.6593\n",
      "Epoch [10/10], Step [220/391], Loss: 2.6609\n",
      "Epoch [10/10], Step [230/391], Loss: 2.6690\n",
      "Epoch [10/10], Step [240/391], Loss: 2.6711\n",
      "Epoch [10/10], Step [250/391], Loss: 2.6709\n",
      "Epoch [10/10], Step [260/391], Loss: 2.6702\n",
      "Epoch [10/10], Step [270/391], Loss: 2.6731\n",
      "Epoch [10/10], Step [280/391], Loss: 2.6763\n",
      "Epoch [10/10], Step [290/391], Loss: 2.6749\n",
      "Epoch [10/10], Step [300/391], Loss: 2.6739\n",
      "Epoch [10/10], Step [310/391], Loss: 2.6713\n",
      "Epoch [10/10], Step [320/391], Loss: 2.6718\n",
      "Epoch [10/10], Step [330/391], Loss: 2.6724\n",
      "Epoch [10/10], Step [340/391], Loss: 2.6705\n",
      "Epoch [10/10], Step [350/391], Loss: 2.6699\n",
      "Epoch [10/10], Step [360/391], Loss: 2.6699\n",
      "Epoch [10/10], Step [370/391], Loss: 2.6688\n",
      "Epoch [10/10], Step [380/391], Loss: 2.6678\n",
      "Epoch [10/10], Step [390/391], Loss: 2.6673\n",
      "Epoch [10/10], Test Accuracy: 31.08%\n"
     ]
    }
   ],
   "source": [
    "# Training loop\n",
    "total_steps = len(train_loader)\n",
    "for epoch in range(num_epochs):\n",
    "    train_loss = 0.0\n",
    "    train_total = 0\n",
    "    train_correct = 0\n",
    "\n",
    "    for i, (images, labels) in enumerate(train_loader):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        # Forward pass\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        # Backward and optimize\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        train_total += labels.size(0)\n",
    "        train_correct += (predicted == labels).sum().item()\n",
    "        # Print training progress\n",
    "        if (i + 1) % 10 == 0:\n",
    "            print(\n",
    "                f\"Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{total_steps}], Loss: {train_loss / (i + 1):.4f}\")\n",
    "\n",
    "    train_acc = 100.0 * train_correct / train_total\n",
    "    train_losses.append(train_loss / total_steps)\n",
    "    train_accs.append(train_acc)\n",
    "    \n",
    "    \n",
    "    # Evaluation on the test set\n",
    "    model.eval()\n",
    "    test_loss = 0.0\n",
    "    test_total = 0\n",
    "    test_correct = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            test_loss += loss.item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            test_total += labels.size(0)\n",
    "            test_correct += (predicted == labels).sum().item()\n",
    "\n",
    "    test_acc = 100.0 * test_correct / test_total\n",
    "    test_losses.append(test_loss / len(test_loader))\n",
    "    test_accs.append(test_acc)\n",
    "\n",
    "    # Print test accuracy for the current epoch\n",
    "    print(f\"Epoch [{epoch+1}/{num_epochs}], Test Accuracy: {test_acc:.2f}%\")\n",
    "\n",
    "    # Switch back to training mode\n",
    "    model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "da551fa6",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'torch' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [1], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Save the model\u001b[39;00m\n\u001b[1;32m      2\u001b[0m PATH \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m./cifar_coordinate_atttion.pth\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m----> 3\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241m.\u001b[39msave(model\u001b[38;5;241m.\u001b[39mstate_dict(), PATH)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'torch' is not defined"
     ]
    }
   ],
   "source": [
    "# Save the model\n",
    "PATH = './cifar_coordinate_atttion.pth'\n",
    "torch.save(model.state_dict(), PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c2df4ca",
   "metadata": {},
   "source": [
    "## 5. Test the network on the test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41c04eae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation on the test set\n",
    "model.eval()\n",
    "test_loss = 0.0\n",
    "test_total = 0\n",
    "test_correct = 0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        # Forward pass\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        test_loss += loss.item()\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        test_total += labels.size(0)\n",
    "        test_correct += (predicted == labels).sum().item()\n",
    "\n",
    "test_acc = 100.0 * test_correct / test_total\n",
    "test_losses.append(test_loss / len(test_loader))\n",
    "test_accs.append(test_acc)\n",
    "\n",
    "# Print test accuracy for the current epoch\n",
    "print(f\"Test Accuracy of the model on the {len(test_loader.dataset)} test images: {test_acc:.2f}%\")\n",
    "\n",
    "# Switch back to training mode\n",
    "model.train()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6197b53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the training and testing losses\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.plot(range(1, num_epochs + 1), train_losses, label='Training Loss')\n",
    "plt.plot(range(1, num_epochs + 1), test_losses, label='Testing Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training and Testing Loss')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# Plot the training and testing acc\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.plot(range(1, num_epochs + 1), train_accs, label='Training Accuracy')\n",
    "plt.plot(range(1, num_epochs + 1), test_accs, label='Testing Accuracy')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.title('Training and Testing Accuracy')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a810aea9-7d2a-4916-88e4-3c98a92160aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
